package org.dizena.blog.starter.config.demo1;

import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.dizena.blog.starter.common.HttpServletUtil;
import org.springframework.core.task.TaskDecorator;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.util.Map;

@Slf4j
public class TransmittableDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        RequestAttributes ras = RequestContextHolder.getRequestAttributes();
        if (ras != null && ras instanceof ServletRequestAttributes) {
            ServletRequestAttributes requestAttributes = (ServletRequestAttributes) ras;
            HttpServletRequest request = requestAttributes.getRequest();
            Map<String, String> header = HttpServletUtil.getHeaderMap(request);
            for (Map.Entry<String, String> e : header.entrySet()) {
                TransmittableUtil.set(e.getKey(), e.getValue());
            }
        }
        return () -> {
            try {
                runnable.run();
            } finally {
                TransmittableUtil.remove();
            }
        };
    }
}